#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
obfus_renamer.py
Rename files in a folder to strongly randomized, Windows-safe names and write a mapping file.
Restore original names later using the mapping.

Features:
- Windows-safe random names (avoid reserved names, invalid chars, trailing dots/spaces).
- Preserve/strip/randomize extension (default: keep original extension).
- Recursive or flat folder processing.
- Include/Exclude glob patterns.
- Dry-run mode with a diff-like preview.
- Collision-safe renaming with retries.
- Optional hashing (SHA-256) and file size/mtime capture in the mapping for integrity checks.
- Optional HMAC signature for the mapping file (detect tampering) via --sign-key.
- Optional keep or randomize timestamps (created/modified) for obfuscation.
- Robust restore that validates paths and detects conflicts.
- Mapping stored with relative paths for portability.
- Optional renaming of directories as well (advanced; disabled by default).

Usage examples:
  Rename (recursive), keep extensions, write mapping.csv:
    python obfus_renamer.py rename --dir "C:\\data" --map mapping.csv --recursive

  Rename only images, exclude .ini files, dry-run first:
    python obfus_renamer.py rename --dir "C:\\data" --map map.csv --recursive \
      --include "*.jpg" --include "*.png" --exclude "*.ini" --dry-run

  Restore from mapping:
    python obfus_renamer.py restore --dir "C:\\data" --map mapping.csv

  Sign the mapping (write HMAC):
    python obfus_renamer.py rename --dir "C:\\data" --map map.csv --sign-key "secret123"

  Verify mapping signature later:
    python obfus_renamer.py verify --map map.csv --sign-key "secret123"
"""

import os
import re
import csv
import sys
import hmac
import time
import fnmatch
import hashlib
import random
import secrets
import string
import argparse
from pathlib import Path
from datetime import datetime

INVALID_WIN_CHARS = set('<>:"/\\|?*')
RESERVED_BASENAMES = {
    "CON", "PRN", "AUX", "NUL",
    *(f"COM{i}" for i in range(1,10)),
    *(f"LPT{i}" for i in range(1,10)),
}

def windows_safe_basename(name: str) -> bool:
    """Check if a basename is Windows-safe (rough, but strict)."""
    if not name or name.strip() != name:
        return False
    if name.endswith('.') or name.endswith(' '):
        return False
    if any(ch in INVALID_WIN_CHARS for ch in name):
        return False
    base_no_ext = name.split('.')[0].upper()
    if base_no_ext in RESERVED_BASENAMES:
        return False
    return True

def random_token(length: int, charset: str, case: str) -> str:
    """Generate a random token from the chosen charset and case."""
    if charset == "hex":
        alphabet = "0123456789abcdef"
    elif charset == "base32":
        alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567"
    elif charset == "url":
        # URL-safe base64 variant, no padding
        alphabet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"
    else:
        # alnum (default)
        alphabet = string.ascii_letters + string.digits

    token = "".join(secrets.choice(alphabet) for _ in range(length))
    if case == "lower":
        token = token.lower()
    elif case == "upper":
        token = token.upper()
    # mixed -> leave as is
    return token

def make_complex_name(minlen: int, maxlen: int, charset: str, case: str, chunks: int) -> str:
    """Build a complex-looking random name with separators to increase 'complexity'."""
    if chunks < 1:
        chunks = 1
    total_len = secrets.choice(range(minlen, maxlen+1))
    # Divide length across chunks
    base_per = max(1, total_len // chunks)
    remainder = total_len - base_per * chunks
    lens = [base_per] * chunks
    for i in range(remainder):
        lens[i % chunks] += 1
    parts = [random_token(l, charset, case) for l in lens]
    # Use separators safe on Windows
    seps = ['-', '_']
    name = parts[0]
    for p in parts[1:]:
        name += secrets.choice(seps) + p
    return name

def ensure_not_reserved(basename: str) -> str:
    """If the name (without extension) is a reserved device name, tweak it."""
    stem = basename.split('.')[0]
    ext = basename[len(stem):]  # includes '.' and rest
    if stem.upper() in RESERVED_BASENAMES:
        stem += "_"
    return stem + ext

def is_relative_to(path: Path, base: Path) -> bool:
    """Compat Path.is_relative_to for Python <3.9"""
    try:
        path.relative_to(base)
        return True
    except Exception:
        return False

def generate_unique_name(target_dir: Path, ext_mode: str, original_ext: str,
                         minlen: int, maxlen: int, charset: str, case: str, chunks: int,
                         existing: set, max_retries: int = 100) -> str:
    """Generate a unique, Windows-safe basename not colliding with existing names in target_dir."""
    for _ in range(max_retries):
        base = make_complex_name(minlen, maxlen, charset, case, chunks)
        base = base.strip(" .")  # avoid trailing dot/space
        if not base:
            continue
        base = ensure_not_reserved(base)
        if ext_mode == "keep":
            new_basename = base + original_ext
        elif ext_mode == "strip":
            new_basename = base
        elif ext_mode == "random":
            # random 2-4 letter extension from alnum
            ext_len = secrets.choice([2,3,4])
            r_ext = ''.join(secrets.choice(string.ascii_lowercase + string.digits) for _ in range(ext_len))
            new_basename = f"{base}.{r_ext}"
        else:
            new_basename = base + original_ext

        if not windows_safe_basename(new_basename):
            continue

        if new_basename not in existing and not (target_dir / new_basename).exists():
            return new_basename
    raise RuntimeError("Failed to generate a unique name after many attempts. Try different parameters.")

def sha256_file(path: Path, chunk_size: int = 1024*1024) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        while True:
            b = f.read(chunk_size)
            if not b:
                break
            h.update(b)
    return h.hexdigest()

def hmac_file(path: Path, key: bytes) -> str:
    h = hmac.new(key, digestmod=hashlib.sha256)
    with path.open("rb") as f:
        while True:
            b = f.read(1024*1024)
            if not b:
                break
            h.update(b)
    return h.hexdigest()

def list_files(base_dir: Path, recursive: bool, include_patterns, exclude_patterns, include_dirs: bool):
    """Yield file (and optionally dir) Paths relative to base_dir according to patterns."""
    def included(rel: Path) -> bool:
        # If no include patterns => include all
        if not include_patterns:
            return True
        s = rel.as_posix()
        return any(fnmatch.fnmatch(s, pat) for pat in include_patterns)

    def excluded(rel: Path) -> bool:
        # If no exclude patterns => exclude none
        if not exclude_patterns:
            return False
        s = rel.as_posix()
        return any(fnmatch.fnmatch(s, pat) for pat in exclude_patterns)

    # Walk
    it = base_dir.rglob("*") if recursive else base_dir.glob("*")

    for path in it:
        rel = path.relative_to(base_dir)
        if path.is_dir():
            if include_dirs and included(rel) and not excluded(rel):
                yield rel
            continue
        if included(rel) and not excluded(rel):
            yield rel

def windows_set_hidden_attribute(path: Path, hidden: bool):
    """Mark file hidden on Windows using attrib command."""
    try:
        if os.name == "nt":
            import subprocess
            if hidden:
                subprocess.run(["attrib", "+H", str(path)], check=False, capture_output=True)
            else:
                subprocess.run(["attrib", "-H", str(path)], check=False, capture_output=True)
    except Exception:
        # Best-effort only
        pass

def set_file_times(path: Path, atime: float, mtime: float):
    try:
        os.utime(path, (atime, mtime))
    except Exception:
        pass

def rename_one(base_dir: Path, rel_path: Path, params, existing_names: set):
    src_abs = base_dir / rel_path
    parent = src_abs.parent
    original_ext = src_abs.suffix if params.ext_mode in ("keep", "random") else ""
    new_name = generate_unique_name(
        parent, params.ext_mode, original_ext,
        params.minlen, params.maxlen, params.charset, params.case, params.chunks,
        existing_names
    )
    dst_abs = parent / new_name
    new_rel = dst_abs.relative_to(base_dir)

    # timestamps
    try:
        stat = src_abs.stat()
        orig_size = stat.st_size
        orig_mtime = stat.st_mtime
        orig_atime = stat.st_atime
    except FileNotFoundError:
        raise FileNotFoundError(f"Source disappeared: {src_abs}")

    if params.dry_run:
        return {
            "old": rel_path.as_posix(),
            "new": new_rel.as_posix(),
            "size": orig_size,
            "mtime": int(orig_mtime),
            "sha256": None,
            "renamed": False
        }

    # Optional hashing
    file_hash = None
    if params.hash:
        try:
            file_hash = sha256_file(src_abs)
        except Exception as e:
            file_hash = f"error:{e}"

    # Do rename
    os.replace(src_abs, dst_abs)

    # Optional timestamp handling
    if params.keep_times:
        set_file_times(dst_abs, orig_atime, orig_mtime)
    elif params.randomize_times:
        now = time.time()
        # Randomize within past 5 years
        five_years = 60*60*24*365*5
        rnd_time = now - secrets.randbelow(five_years)
        set_file_times(dst_abs, rnd_time, rnd_time)

    # Hidden attribute
    if params.set_hidden:
        windows_set_hidden_attribute(dst_abs, True)

    existing_names.add(new_name)

    return {
        "old": rel_path.as_posix(),
        "new": new_rel.as_posix(),
        "size": orig_size,
        "mtime": int(orig_mtime),
        "sha256": file_hash,
        "renamed": True
    }

def read_mapping(map_path: Path):
    # Read raw text, separate out #hmac lines if present
    content = map_path.read_text(encoding="utf-8")
    lines = content.splitlines()
    expected_hmac = None
    if len(lines) >= 2 and lines[-2].startswith("#hmac") and lines[-1].startswith("#hmac"):
        # The last line should be "#hmac,<digest>"
        parts = lines[-1].split(",", 1)
        if len(parts) == 2:
            expected_hmac = parts[1].strip()
        # Strip the last two lines for CSV parsing
        lines = lines[:-2]
    # Parse CSV on the remaining lines
    from io import StringIO
    reader = csv.DictReader(StringIO("\n".join(lines)))
    rows = []
    for row in reader:
        if not row:
            continue
        rows.append(row)
    return rows, expected_hmac

def write_mapping(map_path: Path, rows, sign_key):
    fieldnames = ["old_rel", "new_rel", "size", "mtime", "sha256"]
    # Write to temp and replace for safety
    tmp = map_path.with_suffix(map_path.suffix + ".tmp")
    with tmp.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for r in rows:
            writer.writerow({
                "old_rel": r["old"],
                "new_rel": r["new"],
                "size": r.get("size", ""),
                "mtime": r.get("mtime", ""),
                "sha256": r.get("sha256", ""),
            })
    if sign_key:
        # compute HMAC over file bytes
        digest = hmac_file(tmp, sign_key.encode("utf-8"))
        # Append a tiny 2-row section that will be ignored by strict CSV readers (second header)
        with tmp.open("a", newline="", encoding="utf-8") as f:
            f.write("\n#hmac,value\n#hmac,{}\n".format(digest))

    # Atomic replace
    os.replace(tmp, map_path)

def verify_signature(map_path: Path, sign_key: str):
    # Recompute HMAC over everything except the last 2 lines if they start with "#hmac"
    content = map_path.read_text(encoding="utf-8")
    lines = content.splitlines()
    trimmed = lines[:]
    expected = None
    if len(lines) >= 2 and lines[-2].startswith("#hmac") and lines[-1].startswith("#hmac"):
        trimmed = lines[:-2]
        parts = lines[-1].split(",", 1)
        if len(parts) == 2:
            expected = parts[1].strip()
    tmp = map_path.with_suffix(map_path.suffix + ".verifytmp")
    tmp.write_text("\n".join(trimmed), encoding="utf-8")
    try:
        digest = hmac_file(tmp, sign_key.encode("utf-8"))
        return expected == digest, digest, expected
    finally:
        try:
            tmp.unlink()
        except Exception:
            pass

def do_rename(args):
    base_dir = Path(args.dir).resolve()
    if not base_dir.exists() or not base_dir.is_dir():
        print(f"[ERROR] Directory not found: {base_dir}", file=sys.stderr)
        sys.exit(2)

    map_path = Path(args.map).resolve()
    # Avoid renaming the mapping file itself if it is inside base_dir
    excludes = set(args.exclude or [])
    try:
        if is_relative_to(map_path, base_dir):
            excludes.add(str(map_path.relative_to(base_dir).as_posix()))
    except Exception:
        pass
    # Also avoid renaming this script if it lives inside base_dir
    try:
        script_path = Path(__file__).resolve()
        if is_relative_to(script_path, base_dir):
            excludes.add(str(script_path.relative_to(base_dir).as_posix()))
    except Exception:
        pass

    include_dirs = bool(args.rename_dirs)

    items = list(list_files(base_dir, args.recursive, args.include, list(excludes), include_dirs))
    # Separate dirs and files to avoid renaming directories before files within them
    dir_items = [p for p in items if (base_dir / p).is_dir()]
    file_items = [p for p in items if (base_dir / p).is_file()]

    # For files, ensure deterministic order for nicer diffs
    file_items.sort(key=lambda p: p.as_posix())

    print(f"[INFO] Found {len(file_items)} files" + (f" and {len(dir_items)} directories" if include_dirs else ""))

    existing_names = set()
    results = []

    # Rename files first
    for rel in file_items:
        try:
            res = rename_one(base_dir, rel, args, existing_names)
            results.append(res)
            if args.dry_run:
                print(f"DRY-RUN  {res['old']}  ->  {res['new']}")
            else:
                print(f"RENAMED  {res['old']}  ->  {res['new']}")
        except Exception as e:
            print(f"[ERROR] {rel}: {e}", file=sys.stderr)

    # Optionally rename directories (post-order to avoid moving parents before children)
    if include_dirs:
        dir_items.sort(key=lambda p: len(p.as_posix().split('/')), reverse=True)
        for rel in dir_items:
            src_abs = base_dir / rel
            parent = src_abs.parent
            original_ext = ""  # directories have no extension
            try:
                new_name = generate_unique_name(parent, args.ext_mode, original_ext,
                                                args.minlen, args.maxlen, args.charset, args.case, args.chunks,
                                                existing_names)
                dst_abs = parent / new_name
                new_rel = dst_abs.relative_to(base_dir)
                if args.dry_run:
                    print(f"DRY-RUN  [DIR] {rel.as_posix()}  ->  {new_rel.as_posix()}")
                else:
                    os.replace(src_abs, dst_abs)
                    if args.set_hidden:
                        windows_set_hidden_attribute(dst_abs, True)
                    print(f"RENAMED  [DIR] {rel.as_posix()}  ->  {new_rel.as_posix()}")
            except Exception as e:
                print(f"[ERROR] [DIR] {rel}: {e}", file=sys.stderr)

    # Write mapping for files only
    if not args.dry_run:
        write_mapping(map_path, results, args.sign_key)
        print(f"[OK] Mapping written: {map_path}")
        if args.sign_key:
            ok, digest, expected = verify_signature(map_path, args.sign_key)
            if ok:
                print(f"[OK] Mapping HMAC verified: {digest}")
            else:
                print(f"[WARN] Mapping HMAC mismatch (computed {digest}, expected {expected})")

def do_restore(args):
    base_dir = Path(args.dir).resolve()
    map_path = Path(args.map).resolve()
    if not base_dir.exists() or not base_dir.is_dir():
        print(f"[ERROR] Directory not found: {base_dir}", file=sys.stderr)
        sys.exit(2)
    if not map_path.exists():
        print(f"[ERROR] Mapping file not found: {map_path}", file=sys.stderr)
        sys.exit(2)

    if args.sign_key:
        ok, digest, expected = verify_signature(map_path, args.sign_key)
        if not ok:
            print(f"[ERROR] Mapping HMAC verification failed. (computed {digest}, expected {expected})", file=sys.stderr)
            sys.exit(3)
        else:
            print(f"[OK] Mapping HMAC verified.")

    rows, _ = read_mapping(map_path)
    if not rows:
        print("[ERROR] Mapping appears empty.", file=sys.stderr)
        sys.exit(2)

    # Plan operations
    ops = []
    for row in rows:
        old_rel = Path(row["old_rel"])
        new_rel = Path(row["new_rel"])
        src_abs = base_dir / new_rel
        dst_abs = base_dir / old_rel
        ops.append((src_abs, dst_abs, row))

    # Detect conflicts (destination already exists but is not the same file)
    conflicts = []
    for src_abs, dst_abs, row in ops:
        if not src_abs.exists():
            conflicts.append((src_abs, "missing source"))
        elif dst_abs.exists() and src_abs.resolve() != dst_abs.resolve():
            conflicts.append((dst_abs, "destination exists"))

    if conflicts:
        print("[ERROR] Conflicts detected; resolve manually or use a different directory:", file=sys.stderr)
        for p, reason in conflicts[:20]:
            print(f"  - {p} ({reason})", file=sys.stderr)
        if len(conflicts) > 20:
            print(f"  ... and {len(conflicts)-20} more", file=sys.stderr)
        sys.exit(4)

    # Perform
    for src_abs, dst_abs, row in ops:
        if args.dry_run:
            print(f"DRY-RUN  {src_abs.relative_to(base_dir)}  ->  {dst_abs.relative_to(base_dir)}")
            continue
        dst_abs.parent.mkdir(parents=True, exist_ok=True)
        os.replace(src_abs, dst_abs)
        print(f"RESTORED {src_abs.relative_to(base_dir)}  ->  {dst_abs.relative_to(base_dir)}")

    print("[OK] Restore completed.")

def do_verify(args):
    map_path = Path(args.map).resolve()
    if not map_path.exists():
        print(f"[ERROR] Mapping file not found: {map_path}", file=sys.stderr)
        sys.exit(2)
    if not args.sign_key:
        print("[ERROR] --sign-key is required for verify.", file=sys.stderr)
        sys.exit(2)
    ok, digest, expected = verify_signature(map_path, args.sign_key)
    if ok:
        print(f"[OK] Mapping HMAC verified: {digest}")
    else:
        print(f"[ERROR] Mapping HMAC mismatch. (computed {digest}, expected {expected})", file=sys.stderr)
        sys.exit(3)

def build_arg_parser():
    p = argparse.ArgumentParser(description="Obfuscating file renamer with reversible mapping.")
    sub = p.add_subparsers(dest="cmd", required=True)

    # rename
    pr = sub.add_parser("rename", help="Rename files and write mapping.")
    pr.add_argument("--dir", required=True, help="Target directory to process.")
    pr.add_argument("--map", required=True, help="Path to mapping CSV to create/overwrite.")
    pr.add_argument("--recursive", action="store_true", help="Process subdirectories.")
    pr.add_argument("--include", action="append", help="Glob to include (relative to base). Repeatable.", default=None)
    pr.add_argument("--exclude", action="append", help="Glob to exclude (relative to base). Repeatable.", default=None)
    pr.add_argument("--ext-mode", choices=["keep","strip","random"], default="keep", help="Extension handling.")
    pr.add_argument("--minlen", type=int, default=18, help="Minimum length of random base name.")
    pr.add_argument("--maxlen", type=int, default=30, help="Maximum length of random base name.")
    pr.add_argument("--charset", choices=["alnum","hex","base32","url"], default="alnum", help="Character set for names.")
    pr.add_argument("--case", choices=["mixed","lower","upper"], default="mixed", help="Letter casing.")
    pr.add_argument("--chunks", type=int, default=3, help="Split name into N chunks with separators for complexity.")
    pr.add_argument("--hash", action="store_true", help="Compute SHA-256 for each file into the mapping (slower).")
    pr.add_argument("--dry-run", action="store_true", help="Preview without renaming or writing mapping.")
    pr.add_argument("--keep-times", action="store_true", help="Preserve original atime/mtime on renamed files.")
    pr.add_argument("--randomize-times", action="store_true", help="Randomize atime/mtime within the last ~5 years.")
    pr.add_argument("--set-hidden", action="store_true", help="Set Windows Hidden attribute on renamed items.")
    pr.add_argument("--rename-dirs", action="store_true", help="Also rename directories (advanced).")
    pr.add_argument("--sign-key", help="HMAC-SHA256 key to sign the mapping for tamper detection.")

    # restore
    ps = sub.add_parser("restore", help="Restore original names using mapping.")
    ps.add_argument("--dir", required=True, help="Base directory to restore within.")
    ps.add_argument("--map", required=True, help="Mapping CSV to read.")
    ps.add_argument("--dry-run", action="store_true", help="Preview restore without making changes.")
    ps.add_argument("--sign-key", help="HMAC key to verify mapping integrity before restoring.")

    # verify mapping signature
    pv = sub.add_parser("verify", help="Verify HMAC signature of a mapping file.")
    pv.add_argument("--map", required=True, help="Mapping CSV to verify.")
    pv.add_argument("--sign-key", required=True, help="HMAC key used when creating the mapping.")

    return p

def main():
    parser = build_arg_parser()
    args = parser.parse_args()

    # sanity checks
    if args.cmd == "rename" and args.keep_times and args.randomize_times:
        print("[ERROR] Choose either --keep-times or --randomize-times, not both.", file=sys.stderr)
        sys.exit(2)

    if args.cmd == "rename":
        do_rename(args)
    elif args.cmd == "restore":
        do_restore(args)
    elif args.cmd == "verify":
        do_verify(args)
    else:
        parser.print_help()

if __name__ == "__main__":
    main()
